from PIL import Image
from matplotlib import pyplot as plt
from torchvision.transforms import transforms

from utils import to_abs_path


def do():
    path = to_abs_path("./data/flower_data/train/roses/12240303_80d87f77a3_n.jpg")
    img = Image.open(path)
    # 打印图像尺寸
    print(img.size)
    plt.imshow(img)
    plt.show()

    # 水平翻转
    tfm_hf = transforms.RandomHorizontalFlip(p=1.0)
    t_img = tfm_hf(img)
    plt.imshow(t_img)
    plt.show()

    # 垂直翻转
    tfm_vf = transforms.RandomVerticalFlip(p=1.0)
    t_img = tfm_vf(img)
    plt.imshow(t_img)
    plt.show()

    # 随机裁剪
    tfm_rc = transforms.RandomCrop(size=(128, 128))
    t_img = tfm_rc(img)
    plt.imshow(t_img)
    plt.show()

    # 随机旋转
    tfm_rot = transforms.RandomRotation(degrees=45)
    t_img = tfm_rot(img)
    plt.imshow(t_img)
    plt.show()

    # 尺寸调整
    tfm_rs = transforms.Resize(size=(100, 100))
    t_img = tfm_rs(img)
    plt.imshow(t_img)
    plt.show()

    # 颜色变换
    # 亮度、对比度、饱和度参数范围：[0.0, 1.0]，随机变化为原图的 [1-param, 1+param]
    # 色调参数范围：[-0.5, 0.5],随机变化为原图的 [-param, +param]
    tfm_color = transforms.ColorJitter(brightness=0.5, contrast=0.5, saturation=0.5, hue=0.5)
    t_img = tfm_color(img)
    plt.imshow(t_img)
    plt.show()

    # 转为灰度图
    tfm_gray = transforms.RandomGrayscale(p=1.0)
    t_img = tfm_gray(img)
    plt.imshow(t_img)
    plt.show()

    # 组合增强
    tfm_com = transforms.Compose([
        tfm_hf, tfm_vf, tfm_rc, tfm_rot, tfm_rs, tfm_color, tfm_gray
    ])
    t_img = tfm_com(img)
    plt.imshow(t_img)
    plt.show()
    pass


if __name__ == '__main__':
    do()
